#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

import torch
import numpy as np
import os


def token_shift(x: torch.Tensor, rwkvag: torch.Tensor, h0: torch.Tensor) -> torch.Tensor:
    """
    通道混合函数。

    Args:
        x (torch.Tensor): 输入张量，形状为[Batch, Seq_length, N_embd]。
        rwkvag (torch.Tensor): 输入张量(模型权重），形状为[6, 1, 1, N_embd]。
        h0 (torch.Tensor): 输入张量(state状态)，形状为[Batch, 1, N_embd]。
    Returns:
        xr,xw, xk, xv, xa, xg (torch.Tensor):输出r, w, k, v, a, g混合后的张量,形状均为[Batch, Seq_length, N_embd]
        ht (torch.Tensor): 输出张量(state状态),形状为[Batch, 1, N_embd]。
    """
    batch_size, seq_length = x.shape[0], x.shape[1]
    if seq_length == 1:
        sx = (h0 - x)
        ht = x
    else:
        h0 = h0.view(batch_size,1,-1)
        xx = torch.cat([h0, x[:, :-1, :]], dim=1)
        sx = (xx - x)
        ht = x[:, -1, :]
    xr = x + rwkvag[0] * sx
    xw = x + rwkvag[1] * sx
    xk = x + rwkvag[2] * sx
    xv = x + rwkvag[3] * sx
    xa = x + rwkvag[4] * sx
    xg = x + rwkvag[5] * sx
    return xr, xw, xk, xv, xa, xg, ht

if __name__ == '__main__':
    B = 17
    T = 1
    Nembd = 2560

    x = np.random.uniform(0, 100, [B, T, Nembd]).astype(np.float16)
    rwkvag = np.random.uniform(0, 1, [6, 1, 1, Nembd]).astype(np.float16)
    h0 = np.random.uniform(0, 5, [B, 1, Nembd]).astype(np.float16)

    xr, xw, xk, xv, xa, xg, ht = token_shift(torch.from_numpy(x), torch.from_numpy(rwkvag), torch.from_numpy(h0))

    print(xr.shape)
    print(ht.shape)

    xr_golden = xr.numpy()
    xw_golden = xw.numpy()
    xk_golden = xk.numpy()
    xv_golden = xv.numpy()
    xa_golden = xa.numpy()
    xg_golden = xg.numpy()
    ht_golden = ht.numpy()

    os.system("mkdir -p input")
    os.system("mkdir -p output")
    
    x.tofile("./input/input_x.bin")
    rwkvag.tofile("./input/input_rwkvag.bin")
    h0.tofile("./input/input_h0.bin")
    xr_golden.tofile("./output/golden_xr.bin")
    xw_golden.tofile("./output/golden_xw.bin")
    xk_golden.tofile("./output/golden_xk.bin")
    xv_golden.tofile("./output/golden_xv.bin")
    xa_golden.tofile("./output/golden_xa.bin")
    xg_golden.tofile("./output/golden_xg.bin")
    ht_golden.tofile("./output/golden_ht.bin")



